import pandas as pd
import numpy as np
import os
import talib
from rqalpha.my_factors.base import Factor


class BookImblance(Factor):

    def run(self, sys_freq):
        data = self.source_data
        # print(data)
        _key = list(data.keys())[0]
        index = data[_key].index
        # 第一个因子
        factor1 = data['bidv'].sum(axis=1) / data['askv'].sum(axis=1)
        print('factor1:', '\n', factor1)
        mid_price = (data['askp'][0] + data['bidp'][0]) / 2
        # 将sys_freq和sample_freq全部转化为s级别
        sys_freq = float(sys_freq.split('ms')[0])/1000
        # sample_freq = float(sample_freq.split('s')[0])
        print('sys freq:', sys_freq)
        # print('sample freq:', sample_freq)
        factor_df = pd.DataFrame(data=factor1, index=index, columns=['factor1'])
        # print(factor_df)
        # data[_key].to_csv(r'C:\Users\huajia\Desktop\rqalpha3\rqalpha\data_temp.csv')
        # window = int(sample_freq / freq)
        # print('window is:', window, 'a.shape is:', a.shape[0])
        # cal_idx, window = self.get_index1(data[_key], sys_freq, sample_freq)
        # ret = self.ic_calc(cal_idx, window, factor, mid_price)
        # factor_df = pd.DataFrame(data=ret, index=index)
        # print(factor_df)
        # 第二个因子
        factor2 = data['bidv'].sum(axis=1).pct_change(periods=10) / data['askv'].sum(axis=1).pct_change(periods=10)
        print('factor2:', '\n', factor2)
        factor_df['factor2'] = factor2
        # 第三个因子
        factor3 = talib.EMA(factor1.values, timeperiod=20)
        print('factor3:', '\n', factor3)
        factor_df['factor3'] = factor3
        # 第四个因子
        factor4 = talib.EMA(factor2.values, timeperiod=20)
        print('factor4:', '\n', factor4)
        factor_df['factor4'] = factor4
        # 第五个因子
        total_vol = data['s_volume']['s_vol'] + data['b_volume']['b_vol']
        idx = self.get_index1(total_vol, 20)
        avg_price = self.get_avg_price(idx, data['price'])
        vwap_price = self.get_vwap_price(idx, total_vol, avg_price)
        factor5 = vwap_price / avg_price
        # print('shape:', vwap_price.shape, avg_price.shape)
        print('factor5:', '\n', factor5)
        factor_df['factor5'] = factor5
        # 第六个因子
        factor6 = talib.EMA(data['b_volume']['b_vol'].values, timeperiod=20) / \
                  talib.EMA(data['s_volume']['s_vol'].values, timeperiod=20)
        print('factor6:', '\n', factor6)
        factor_df['factor6'] = factor6
        # 第七个因子
        factor7 = talib.EMA(data['b_unit']['b_unit'].values, timeperiod=20) / \
                  talib.EMA(data['s_unit']['s_unit'].values, timeperiod=20)
        print('factor7:', '\n', factor7)
        factor_df['factor7'] = factor7

        # 回归
        # 获取价格涨跌幅
        price_pct_change = data['price']['last'].pct_change()
        # 20分钟
        window = int(1200/0.5)
        idx = self.get_index1(price_pct_change, window)
        weight_li = self.get_weight_bp(idx, factor_df, price_pct_change.values)
        print(weight_li[0:9])
        weight_df = pd.DataFrame(data=weight_li, index=index)
        # print('weight_df:', weight_df)


